import numpy as np
from keras.models import load_model
import matplotlib.pyplot as plt
from loss_functions import cross_entropy_balanced, pixel_error
import argparse
import cv2
from utils import nms
import time
from yolo import YOLO
from vm_test import someProcessDetailed, tellw_stereoMatch2, surf_stereoMatch, tellw_stereoMatch
from configparser import ConfigParser
from volumeMeasure import VolumeMeasure

def args_parse():
    # construct the argument parse and parse the arguments
    ap = argparse.ArgumentParser(description='Keras Training')
    # ========= paths for training
    ap.add_argument("-npath", "--npy_path", required=True,
                    help="path to npy. files to train")
    ap.add_argument("-mpath", "--model_path", required=True,
                    help="path to save the output model")
    ap.add_argument("-name","--model_name", required=True,
                    help="output of model name")
    ap.add_argument("-r", "--rows", required=True, type=int, default=320,
                    help="shape of rows of input image")
    ap.add_argument("-c", "--cols", required=True, type=int, default=480,
                    help="shape of cols of input image")
    args = vars(ap.parse_args())
    return args

def test(args):
    '''
    X_train = np.load(args["npy_path"] + 'X_train_ori.npy')
    X_test = np.load(args["npy_path"] + 'X_test_ori.npy')
    X_val = np.load(args["npy_path"] + 'X_val_ori.npy')
    y_train = np.load(args["npy_path"] + 'y_train_concat.npy')
    y_test = np.load(args["npy_path"] + 'y_test_concat.npy')
    y_val = np.load(args["npy_path"] + 'y_val_concat.npy')
    '''
    X_train = np.load(args["npy_path"] + 'X_train.npy')
    X_test = np.load(args["npy_path"] + 'X_test.npy')
    X_val = np.load(args["npy_path"] + 'X_val.npy')
    y_train = np.load(args["npy_path"] + 'y_train.npy')
    y_test = np.load(args["npy_path"] + 'y_test.npy')
    y_val = np.load(args["npy_path"] + 'y_val.npy')
    model = load_model(args["model_path"] + args["model_name"],
                       custom_objects={'cross_entropy_balanced': cross_entropy_balanced, 'pixel_error': pixel_error})
    # test all images from test.npy
    print(len(X_train))
    for i in range(200):
        y_pred = model.predict(X_train[i].reshape((-1, 320, 480, 3)))[-1]
        '''
        y_pred = y_pred.reshape((320, 480))
        plt.figure(figsize=(25, 16))
        plt.subplot(1, 3, 1)
        plt.imshow(X_train[i], cmap='binary')
        plt.subplot(1, 3, 2)
        plt.imshow(y_train[i].reshape((320, 480)), cmap='binary')
        plt.subplot(1, 3, 3)
        plt.imshow(y_pred, cmap='binary')
        name = str(i) + '.jpg'
        plt.savefig(name)
        '''
        y_pred = model.predict(X_train[i].reshape((-1, 320, 480, 3)))[-1]
        y_pred = y_pred.reshape((320, 480))
        plt.figure(figsize=(25, 16))
        plt.subplot(1, 3, 1)
        plt.imshow(X_train[i]/255, cmap='binary')
        plt.subplot(1, 3, 2)
        plt.imshow(y_train[i].reshape((320, 480)), cmap='binary')
        plt.subplot(1, 3, 3)
        plt.imshow(y_pred, cmap='binary')
        plt.show()

if __name__ == "__main__":
    vm = VolumeMeasure()
    # yolo = YOLO()
    # for i in range(37):
    #     img = cv2.imread('C:/Users/tellw/Desktop/bishe/items/calibed_pics/20200326101821left/%d.jpg'%i)
    #     yolo.detect_image(img, True)
    cfg = ConfigParser()
    cfg.read('SWQT.ini')
    img = cv2.imread('C:/Users/tellw/Desktop/bishe/items/calibed_pics/20200326101821left/13.jpg')
    img = vm.calibrator.rectify_left(img)
    time1 = []
    time2 = []
    time3 = []
    for i in range(25):
        left = cv2.imread('C:/Users/tellw/Desktop/bishe/items/calibed_pics/20200326101821left/%d.jpg' % i, 0)
        right = cv2.imread('C:/Users/tellw/Desktop/bishe/items/calibed_pics/20200326101821right/%d.jpg' % i, 0)
        rec_left = vm.calibrator.rectify_left(left)
        rec_right = vm.calibrator.rectify_right(right)
        rec_left_roi_src = rec_left[25:455, 25:615]
        rec_right_roi_src = rec_right[25:455, 25:615]
        rec_left_roi = someProcessDetailed(rec_left_roi_src, cfg=cfg, bGaussianBlur=1, bCanny=1)
        rec_right_roi = someProcessDetailed(rec_right_roi_src, cfg=cfg, bGaussianBlur=1, bCanny=1)
        s1=time.time()
        tellw_stereoMatch2(rec_left_roi, rec_right_roi, bPointMatched=True)
        s2=time.time()
        surf_stereoMatch(rec_left_roi, rec_right_roi)
        s3= time.time()
        tellw_stereoMatch(rec_left_roi, rec_right_roi, rec_left_roi_src, rec_right_roi_src)
        s4 = time.time()
        time1.append(s2-s1)
        time2.append(s3-s2)
        time3.append(s4-s3)
    x=list(range(25))

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False
    plt.plot(x, time1, linestyle='--', color='k', label='基于轮廓外接矩形框的立体匹配')
    plt.plot(x, time2, linestyle='-', color='k', label='基于SURF的立体匹配')
    plt.plot(x, time3, linestyle='-.', color='k', label='基于互相关运算的立体匹配')
    plt.xlabel('实验编号')
    plt.ylabel('时长/s')
    plt.title('算法运行时间')
    plt.legend()
    plt.show()
